{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 简介\n",
    "\n",
    "这是一款名为‘FunctionAgentWithRetrieval检索增强小助手’的演示应用，展示了如何结合[llama_index](https://github.com/run-llama/llama_index.git)的工具和ERNIEBot的function agent来回答用户的专业知识问题。首先，应用通过构建FaissSearchLlamaIndex在知识库中检索相关内容，然后评估检索内容与问题的相关度。如果内容与问题高度相关，大模型将采用检索增强方式回答问题；否则，大模型将调用工具列表中的工具回答用户的问题。这种设计不仅扩展了大模型的专业领域知识，还保持了其在领域知识之外的通用对话能力\n",
    "\n",
    "构建流程如下："
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 1. 安装llama_index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install llama-index"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 2. 导入第三方库\n",
    "\n",
    "主要是在导入一些必要的Python库和模块，以便实现FunctionAgentWithRetrival的功能。\n",
    "+ os: Python的标准库，用于与操作系统进行交互，如读写文件、管理路径等。\n",
    "+ faiss: 用于向量存储的模块，是用于存储和检索经过嵌入处理的文本或图像的向量表示。\n",
    "+ SimpleDirectoryReader: 用于在文件夹中加载数据\n",
    "+ ErnieEmbeddings: 用于文本嵌入的工具，将文本转换为可以在模型中使用的向量表示。\n",
    "+ FunctionAgent: 这个类实现function calling功能的Agent的类，如问答、对话等。\n",
    "+ WholeMemory: 用于存储和管理代理的记忆的类。\n",
    "+ ERNIEBot: 实现ERNIE Bot的主要类，包含了实现对话功能的主要逻辑。\n",
    "+ erniebot: 这是导入erniebot模块的语句，erniebot是一个包含ERNIE Bot实现的主要代码库。\n",
    "+ load_index_from_storage: 用于从存储系统加载索引。\n",
    "+ VectorStoreIndex: 提供了创建索引、添加数据到索引等功能。\n",
    "+ StorageContext: 用于存储节点的实用容器。\n",
    "+ FaissVectorStore: 提供了一种使用 FAISS 库来存储和检索向量的方法。\n",
    "+ ServiceContext: 服务上下文容器是 LlamaIndex 的实用程序容器。\n",
    "\n",
    "这里只调用llama_index的检索功能，os.environ[\"OPENAI_API_KEY\"]可以赋值空字符\n",
    "\n",
    "这里使用EB_AGENT_ACCESS_TOKEN, 申请地址请参考[accessToken](https://aistudio.baidu.com/index/accessToken)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "os.environ[\"OPENAI_API_KEY\"] = \" \"\n",
    "os.environ[\"EB_AGENT_ACCESS_TOKEN\"] = \"your access token\"\n",
    "\n",
    "from erniebot_agent.agents.function_agent_with_retrieval import FunctionAgentWithRetrieval\n",
    "from erniebot_agent.memory.whole_memory import WholeMemory\n",
    "from erniebot_agent.chat_models.erniebot import ERNIEBot\n",
    "from erniebot_agent.tools import RemoteToolkit\n",
    "from erniebot_agent.extensions.langchain.embeddings import ErnieEmbeddings\n",
    "import faiss\n",
    "from llama_index import (\n",
    "    SimpleDirectoryReader,\n",
    "    load_index_from_storage,\n",
    "    VectorStoreIndex,\n",
    "    StorageContext,\n",
    "    ServiceContext,\n",
    ")\n",
    "from llama_index.vector_stores.faiss import FaissVectorStore"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 3. 预处理\n",
    "\n",
    "## 3.1 下载数据集"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--2024-01-02 02:27:44--  https://paddlenlp.bj.bcebos.com/datasets/examples/construction_regulations.tar\n",
      "Resolving paddlenlp.bj.bcebos.com (paddlenlp.bj.bcebos.com)... 36.110.192.178, 2409:8c04:1001:1002:0:ff:b001:368a\n",
      "Connecting to paddlenlp.bj.bcebos.com (paddlenlp.bj.bcebos.com)|36.110.192.178|:443... connected.\n",
      "HTTP request sent, awaiting response... 200 OK\n",
      "Length: 1239040 (1.2M) [application/x-tar]\n",
      "Saving to: ‘construction_regulations.tar’\n",
      "\n",
      "construction_regula 100%[===================>]   1.18M  3.01MB/s    in 0.4s    \n",
      "\n",
      "2024-01-02 02:27:44 (3.01 MB/s) - ‘construction_regulations.tar’ saved [1239040/1239040]\n",
      "\n",
      "construction_regulations/\n",
      "construction_regulations/城市管理执法办法.pdf\n",
      "construction_regulations/建筑工程设计招标投标管理办法.pdf\n",
      "construction_regulations/建筑业企业资质管理规定.pdf\n",
      "construction_regulations/城市照明管理规定.pdf\n",
      "construction_regulations/城市设计管理办法.pdf\n",
      "construction_regulations/建筑工程施工发包与承包计价管理办法.pdf\n",
      "construction_regulations/市政公用设施抗灾设防管理规定.pdf\n"
     ]
    }
   ],
   "source": [
    "! wget https://paddlenlp.bj.bcebos.com/datasets/examples/construction_regulations.tar\n",
    "! tar xvf construction_regulations.tar"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3.2 构建FaissSearch_lama\n",
    "这段代码定义了一个名为 FaissSearchLlamaIndex 的类，它使用 FAISS 库进行文档搜索。这个类使用 FAISS 来创建一个基于文档的向量存储索引，并提供了一个搜索方法来查询索引并返回最相关的文档。这种方法适用于需要快速且高效地处理大量高维向量数据的场景，如文本相似性搜索。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "class FaissSearchLlamaIndex:\n",
    "    def __init__(\n",
    "        self,\n",
    "        faiss_index,\n",
    "        documents,\n",
    "        embed_model,\n",
    "    ):\n",
    "        # 类的初始化方法，接收 FAISS 索引、文档集合和嵌入模型作为参数\n",
    "        self.vector_store = FaissVectorStore(faiss_index=faiss_index)\n",
    "        # 初始化 FaissVectorStore，用于存储和检索向量\n",
    "        storage_context = StorageContext.from_defaults(vector_store=self.vector_store)\n",
    "        # 创建一个storage_context，配置默认值和向量存储\n",
    "        self.service_context = ServiceContext.from_defaults(embed_model=embed_model)\n",
    "        # 创建一个service_context，配置默认值和嵌入模型\n",
    "        self.index = VectorStoreIndex.from_documents(\n",
    "            documents,\n",
    "            storage_context=storage_context,\n",
    "            show_progress=True,\n",
    "            service_context=self.service_context,\n",
    "        )\n",
    "        # 从文档集合创建一个向量存储索引，并显示进度\n",
    "        self.index.storage_context.persist()\n",
    "        # 持久化索引数据到存储\n",
    "\n",
    "    def search(self, query, topk=10, **kwargs):\n",
    "        # 定义搜索方法，接收查询字符串、返回结果数量上限以及其他参数\n",
    "        storage_context = StorageContext.from_defaults(\n",
    "            vector_store=self.vector_store, persist_dir=\"./storage\"\n",
    "        )\n",
    "        # 创建一个新的存储上下文，指定持久化目录\n",
    "        index = load_index_from_storage(\n",
    "            storage_context=storage_context, service_context=self.service_context, top_k=topk\n",
    "        )\n",
    "        # 从存储中加载索引\n",
    "        retriever = index.as_retriever()\n",
    "        # 从索引中创建检索器\n",
    "        nodes = retriever.retrieve(query)\n",
    "        # 使用检索器对查询进行检索\n",
    "        retrieval_results = []\n",
    "        for doc in nodes:\n",
    "            # 遍历检索到的文档\n",
    "            retrieval_results.append(\n",
    "                {\"content\": doc.node.text, \"score\": doc.score, \"title\": doc.node.metadata[\"file_name\"]}\n",
    "            )\n",
    "            # 将文档的内容、评分和标题添加到结果列表中\n",
    "        return retrieval_results  # 返回搜索结果"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 4.FunctionAgentWithRetrieval\n",
    "## 4.1 建索引库\n",
    "\n",
    "这段代码主要功能是用于创建或加载一个FAISS索引来进行文档相似度匹配。首先，创建一个ErnieEmbeddings对象，用于生成文档的嵌入向量。aistudio_access_token是访问AI Studio的令牌，chunk_size是用于嵌入的文档块的数量"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "aistudio_access_token = os.environ.get(\"EB_AGENT_ACCESS_TOKEN\", \"\")\n",
    "embed_model = ErnieEmbeddings(aistudio_access_token=aistudio_access_token, chunk_size=16)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "接下来利用ErnieEmbeddings来抽取向量构建索引。d是ernie-embedding的向量抽取维度,并向SimpleDirectoryReader传入数据存储的目录"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "d = 384\n",
    "faiss_index = faiss.IndexFlatIP(d)\n",
    "documents = SimpleDirectoryReader(\"construction_regulations\").load_data()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "以下代码是使用FaissSearchLlamaIndex进行搜索的一个例子，流程大致如下：\n",
    "1. 创建FaissSearchLlamaIndex对象，传入embed_model、documents、faiss_index。\n",
    "2. 调用FaissSearchLlamaIndex对象的search函数，并传入查询字符串\"城市管理执法主管部门的职责是什么？\"。\n",
    "3. 将搜索结果以格式化的形式进行打印，这里使用了pprint模块进行美化打印。最终的结果存储在变量res中。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/root/qingzhong/anas/envs/py310_openai/lib/python3.10/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n",
      "Parsing nodes: 100%|██████████████████████████████████████████████████████████████████████████| 82/82 [00:00<00:00, 1680.35it/s]\n",
      "Generating embeddings:  12%|████████▎                                                           | 10/82 [00:00<00:03, 18.85it/s]Retrying requests: Attempt 1 ended with: <Future at 0x7f3944884820 state=finished raised RateLimitError>\n",
      "Generating embeddings:  24%|████████████████▌                                                   | 20/82 [00:02<00:09,  6.68it/s]Retrying requests: Attempt 1 ended with: <Future at 0x7f3944885180 state=finished raised RateLimitError>\n",
      "Generating embeddings:  37%|████████████████████████▉                                           | 30/82 [00:04<00:09,  5.75it/s]Retrying requests: Attempt 1 ended with: <Future at 0x7f3944884d00 state=finished raised RateLimitError>\n",
      "Generating embeddings:  49%|█████████████████████████████████▏                                  | 40/82 [00:07<00:08,  5.07it/s]Retrying requests: Attempt 1 ended with: <Future at 0x7f3944885150 state=finished raised RateLimitError>\n",
      "Generating embeddings:  61%|█████████████████████████████████████████▍                          | 50/82 [00:09<00:06,  4.89it/s]Retrying requests: Attempt 1 ended with: <Future at 0x7f3944885090 state=finished raised RateLimitError>\n",
      "Generating embeddings:  73%|█████████████████████████████████████████████████▊                  | 60/82 [00:11<00:04,  4.90it/s]Retrying requests: Attempt 1 ended with: <Future at 0x7f3944885330 state=finished raised RateLimitError>\n",
      "Generating embeddings:  85%|██████████████████████████████████████████████████████████          | 70/82 [00:13<00:02,  4.90it/s]Retrying requests: Attempt 1 ended with: <Future at 0x7f3944884f70 state=finished raised RateLimitError>\n",
      "Generating embeddings:  98%|██████████████████████████████████████████████████████████████████▎ | 80/82 [00:15<00:00,  4.87it/s]Retrying requests: Attempt 1 ended with: <Future at 0x7f3944885660 state=finished raised RateLimitError>\n",
      "Generating embeddings: 100%|████████████████████████████████████████████████████████████████████| 82/82 [00:17<00:00,  4.75it/s]\n",
      "Retrying requests: Attempt 1 ended with: <Future at 0x7f3944902350 state=finished raised RateLimitError>\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[{'content': '住房和城乡建设部规章 \\n'\n",
      "             '   X住房和城乡建设部发布  \\n'\n",
      "             '- 11 -   第三十八条 城市管理执法主管部门应当向社会公布投诉、举报电话及其他监督方式。   '\n",
      "             '城市管理执法主管部门应当为投诉人、举报人保密。   第三十九条 城市管理执法主管部门违反本办法规定， 有下列行为之一的， '\n",
      "             '由上级城市管理执法主管部门或者有关部门责令改正，通报批评；情节严重的，对直接负责的主管人员和其他直接责任人员依法给予处分。   '\n",
      "             '（一）没有法定依据实施行政处罚的；   （二）违反法定程序实施行政处罚的；   （三）以罚款、没收违法所得作为经费来源的；   '\n",
      "             '（四）使用、截留、损毁或者擅自处置查封、扣押物品的；   （五）其他违反法律法规和本办法规定的。   第四十条 '\n",
      "             '非城市管理执法人员着城市管理执法制式服装的，城市管理执法主管部门应当予以纠正，依法追究法律责任。',\n",
      "  'score': 0.5666415095329285,\n",
      "  'title': '城市管理执法办法.pdf'},\n",
      " {'content': '住房和城乡建设部规章 \\n'\n",
      "             '   X住房和城乡建设部发布  \\n'\n",
      "             '- 10 -   第三十四条 城市管理执法主管部门应当通过门户网站、 '\n",
      "             '办事窗口等渠道或者场所，公开行政执法职责、权限、依据、监督方式等行政执法信息。   第六章 协作与配合     第三十五条 '\n",
      "             '城市管理执法主管部门应当与有关部门建立行政执法信息互通共享机制， 及时通报行政执法信息和相关行政管理信息。   第三十六条 '\n",
      "             '城市管理执法主管部门可以对城市管理执法事项实行网格化管理。   第三十七条 '\n",
      "             '城市管理执法主管部门在执法活动中发现依法应当由其他部门查处的违法行为， 应当及时告知或者移送有关部门。   第七章 执法监督',\n",
      "  'score': 0.5620024800300598,\n",
      "  'title': '城市管理执法办法.pdf'}]\n"
     ]
    }
   ],
   "source": [
    "faiss_search = FaissSearchLlamaIndex(embed_model=embed_model, documents=documents, faiss_index=faiss_index)\n",
    "query = \"城市管理执法主管部门的职责是什么？\"\n",
    "res = faiss_search.search(query)\n",
    "from pprint import pprint\n",
    "\n",
    "pprint(res)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4.2 构建FunctionAgentWithRetrieval\n",
    "FunctionAgentWithRetrieval可以决定是否进行检索增强来回答用户问题，以下是其示例。\n",
    "### 4.2.1 检索增强"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 创建一个ERNIEBot实例，使用\"ernie-3.5\"模型。\n",
    "llm = ERNIEBot(model=\"ernie-3.5\")\n",
    "# 创建一个WholeMemory实例。这是一个用于存储对话历史和上下文信息的类，有助于模型理解和持续对话。\n",
    "memory = WholeMemory()\n",
    "# 调用文本语音转换工具。\n",
    "tts_tool = RemoteToolkit.from_aistudio(\"texttospeech\").get_tools()[0]\n",
    "# 创建一个FunctionAgentWithRetrieval实例。这个代理将使用上面创建的ERNIEBot模型、WholeMemory和faiss_search，同时传入了一个名为tts_tool的工具。\n",
    "agent = FunctionAgentWithRetrieval(\n",
    "    llm=llm, tools=[tts_tool], memory=memory, knowledge_base=faiss_search, threshold=0.5\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'role': 'user', 'content': '检索结果:\\n\\n    第1个段落: 住房和城乡建设部规章 \\n   X住房和城乡建设部发布  \\n- 9 -   第三十条 城市管理执法主管部门不得对罚款、 没收违法所得设定任务和目标。   罚款、没收违法所得的款项，应当按照规定全额上缴。   第三十一条 城市管理执法主管部门应当确定法制审核机构，配备一定比例符合条件的法制审核人员，对重大执法决定在执法主体、管辖权限、执法程序、事实认定、法律适用等方面进行法制审核。   第三十二条 城市管理执法主管部门开展执法活动， 应当使用统一格式的行政执法文书。   第三十三条 行政执法文书的送达， 依照民事诉讼法等法律规定执行。   当事人提供送达地址或者同意电子送达的， 可以按照其提供的地址或者传真、电子邮件送达。   采取直接、留置、邮寄、委托、转交等方式无法送达的，可以通过报纸、门户网站等方式公告送达。\\n\\n    第2个段落: 住房和城乡建设部规章 \\n   X住房和城乡建设部发布  \\n- 4 -   第十一条 城市管理执法事项范围确定后，应当向社会公开。   第十二条 城市管理执法主管部门集中行使原由其他部门行使的行政处罚权的，应当与其他部门明确职责权限和工作机制。   第三章 执法主体     第十三条 城市管理执法主管部门按照权责清晰、事权统一、精简效能的原则设置执法队伍。   第十四条 直辖市、 设区的市城市管理执法推行市级执法或者区级执法。   直辖市、设区的市的城市管理执法事项，市辖区人民政府城市管理执法主管部门能够承担的，可以实行区级执法。   直辖市、 设区的市人民政府城市管理执法主管部门可以承担跨区域和重大复杂违法案件的查处。\\n\\n检索语句: 城乡建设部规章中，城市管理执法第三章，第十三条是什么？\\n请根据以上检索结果回答检索语句的问题'}\n",
      "{'role': 'assistant', 'content': '根据提供的检索结果，城乡建设部规章中，城市管理执法第三章，第十三条的内容如下：\\n\\n第十三条 城市管理执法主管部门按照权责清晰、事权统一、精简效能的原则设置执法队伍。', 'function_call': None, 'plugin_info': None, 'search_info': {'results': [{'index': 1, 'url': '', 'title': '城市管理执法办法.pdf'}, {'index': 2, 'url': '', 'title': '城市管理执法办法.pdf'}]}}\n"
     ]
    }
   ],
   "source": [
    "# 定义一个查询字符串，这个查询是关于\"城乡建设部规章中，城市管理执法第三章，第十三条\"的内容。\n",
    "query = \"城乡建设部规章中，城市管理执法第三章，第十三条是什么？\"\n",
    "# 使用agent的async_run方法来异步执行查询。由于这是异步操作，因此需要使用'await'关键字。\n",
    "response = await agent._run(query)\n",
    "messages = response.chat_history\n",
    "for item in messages:\n",
    "    print(item.to_dict())"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 4.2.2 调用tool示例\n",
    "在这个例子中，FunctionAgentWithRetrieval没有利用其检索信息，而是直接调用工具回答用户的问题"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'role': 'user', 'content': '把上一轮的检索内容转换为语音'}\n",
      "{'role': 'assistant', 'content': '', 'function_call': {'name': 'texttospeech/v1.6/tts', 'thoughts': '用户需要将检索内容转换为语音。', 'arguments': '{\"tex\":\"住房和城乡建设部规章 第三章 第十三条 城市管理执法主管部门按照权责清晰、事权统一、精简效能的原则设置执法队伍。\"}'}, 'plugin_info': None, 'search_info': None}\n",
      "{'role': 'function', 'name': 'texttospeech/v1.6/tts', 'content': '{\"audio\": \"file-local-ad9def7e-a916-11ee-aa47-fa280016e4f3\", \"prompt\": \"参考工具说明中对各个结果字段的描述，提取工具调用结果中的信息，生成一段通顺的文本满足用户的需求。请务必确保每个符合\\'file-\\'格式的字段只出现一次，无需将其转换为链接，也无需添加任何HTML、Markdown或其他格式化元素。\"}'}\n",
      "{'role': 'assistant', 'content': '根据你的请求，我已将“住房和城乡建设部规章 第三章 第十三条 城市管理执法主管部门按照权责清晰、事权统一、精简效能的原则设置执法队伍。”这段文字转换为语音，格式为file-local-ad9def7e-a916-11ee-aa47-fa280016e4f3。', 'function_call': None, 'plugin_info': None, 'search_info': None}\n"
     ]
    }
   ],
   "source": [
    "# 定义一个查询字符串，这个查询是关于\"把上一轮的检索内容转换为语音\"的内容。\n",
    "response = await agent._run(\"把上一轮的检索内容转换为语音\")\n",
    "# 使用agent的async_run方法来异步执行查询。由于这是异步操作，因此需要使用'await'关键字。\n",
    "messages = response.chat_history\n",
    "for item in messages:\n",
    "    print(item.to_dict())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "openai",
   "language": "python",
   "name": "py310_openai"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
